import pickle
import pandas as pd
import os


def save_pickle(data, save_name):
    with open(save_name, 'wb') as f:
        pickle.dump(data, f)
    return save_name


def cut_xy(x_se: pd.Series, y_se: pd.Series):
    t_in = [t for t in x_se.index if t in y_se.index]
    x_se = x_se[t_in]
    y_se = y_se[t_in]
    print('x_se from', x_se.index[0], ' to ', x_se.index[-1])
    print('y_se from', y_se.index[0], ' to ', y_se.index[-1])
    return x_se, y_se


def cal_tztd(x_se: pd.Series, y_se: pd.Series, cal_y=False, t_cut='2020-01-01'):
    """
    求两个时间序列同涨同跌比率
    :param cal_y: 计算同涨同跌的时候基准点的选择
    :param x_se:
    :param y_se:
    :return:
    """
    x_se.index = pd.to_datetime(x_se.index)
    y_se.index = pd.to_datetime(y_se.index)
    x_se, y_se = cut_xy(x_se, y_se)
    x_se.index = pd.to_datetime(x_se.index)
    y_se.index = pd.to_datetime(y_se.index)
    x_se = x_se[x_se.index > pd.to_datetime(t_cut)]
    y_se = y_se[y_se.index > pd.to_datetime(t_cut)]
    acc_li = []
    acc_dic = {}
    for i in range(1, len(x_se)):
        t = x_se.index[i]
        if cal_y:
            v_x = (x_se[i] - y_se[i - 1])
        else:
            v_x = (x_se[i] - x_se[i - 1])
        v_y = (y_se[i] - y_se[i - 1])
        v = v_x * v_y
        if v > 0:
            direction = 1
        elif v < 0:
            direction = -1
        else:
            direction = 0
        acc_li.append(direction)
        if i > 1:
            acc_dic[t] = (direction, (t_last, t))
        t_last = t
    acc_li_all = [i for i in acc_li if i in [1, -1]]
    acc_li_1 = [i for i in acc_li if i == 1]
    acc_ = len(acc_li_1) / len(acc_li_all)
    return acc_, acc_dic


def list_excel():
    pths = os.listdir('input/')
    pths = [s for s in pths if 'xlsx' in s]
    for p in pths:
        print(p)
    return pths